# -*- coding: utf-8 -*-
"""
流水线管理类
负责协调和管理整个视频生成流水线
"""

import logging
from pathlib import Path
from typing import Dict, List, Optional, Any, Callable

from .step_base import BaseStep, StepResult, StepStatus
from .checkpoint_manager import CheckpointManager

logger = logging.getLogger(__name__)


class Pipeline:
    """流水线管理类"""

    def __init__(
        self,
        name: str,
        steps: List[BaseStep],
        checkpoint_dir: Optional[Path] = None,
        auto_save: bool = True,
    ):
        """
        初始化流水线

        Args:
            name: 流水线名称
            steps: 步骤列表
            checkpoint_dir: 断点目录
            auto_save: 是否自动保存断点
        """
        self.name = name
        self.steps = steps
        self.auto_save = auto_save
        self.current_step_index = 0
        self.checkpoint_manager = CheckpointManager(
            checkpoint_dir or Path("checkpoints")
        )

        # 进度回调
        self.progress_callback: Optional[Callable[[int, str], None]] = None
        self.step_start_callback: Optional[Callable[[str], None]] = None
        self.step_complete_callback: Optional[Callable[[str, StepResult], None]] = None

    def set_progress_callback(self, callback: Callable[[int, str], None]):
        """
        设置进度回调

        Args:
            callback: 进度回调函数 (percent, message)
        """
        self.progress_callback = callback
        for step in self.steps:
            step.progress_callback = callback

    def set_step_start_callback(self, callback: Callable[[str], None]):
        """
        设置步骤开始回调

        Args:
            callback: 步骤开始回调函数 (step_name)
        """
        self.step_start_callback = callback

    def set_step_complete_callback(self, callback: Callable[[str, StepResult], None]):
        """
        设置步骤完成回调

        Args:
            callback: 步骤完成回调函数 (step_name, result)
        """
        self.step_complete_callback = callback

    def _report_progress(self, percent: int, message: str = ""):
        """报告进度"""
        if self.progress_callback:
            self.progress_callback(percent, message)
        logger.info(f"[{self.name}] {percent}% - {message}")

    def _save_checkpoint(self):
        """保存断点"""
        if not self.auto_save:
            return

        # 收集所有步骤状态
        steps_state = {}
        for step in self.steps:
            steps_state[step.name] = step.get_status_dict()

        # 保存流水线状态
        current_step = (
            self.steps[self.current_step_index].name
            if self.current_step_index < len(self.steps)
            else None
        )

        self.checkpoint_manager.save_pipeline_state(
            pipeline_name=self.name,
            current_step=current_step,
            steps_state=steps_state,
        )

    def run(self, from_checkpoint: bool = False) -> StepResult:
        """
        运行流水线

        Args:
            from_checkpoint: 是否从断点恢复

        Returns:
            执行结果
        """
        logger.info(f"Starting pipeline: {self.name}")

        # 如果需要从断点恢复
        if from_checkpoint:
            state = self.checkpoint_manager.load_pipeline_state()
            if state and state.get("pipeline_name") == self.name:
                self._restore_from_checkpoint(state)
            else:
                logger.warning("No valid checkpoint found, starting fresh")

        # 执行步骤
        try:
            for i, step in enumerate(self.steps):
                self.current_step_index = i

                # 检查是否已跳过
                if step.status == StepStatus.SKIPPED:
                    logger.info(f"Step {i+1}/{len(self.steps)} skipped: {step.name}")
                    self._report_progress(
                        int((i / len(self.steps)) * 100),
                        f"Skipped: {step.name}"
                    )
                    continue

                # 执行步骤
                self._report_progress(
                    int((i / len(self.steps)) * 100),
                    f"Running: {step.name}"
                )

                if self.step_start_callback:
                    self.step_start_callback(step.name)

                result = step.execute()

                if self.step_complete_callback:
                    self.step_complete_callback(step.name, result)

                # 保存断点
                if self.auto_save:
                    self._save_checkpoint()

                # 检查结果
                if not result.success:
                    self._report_progress(
                        int((i / len(self.steps)) * 100),
                        f"Failed: {step.name} - {result.message}"
                    )
                    return result

            # 所有步骤完成
            self._report_progress(100, "Pipeline completed successfully")

            return StepResult(
                success=True,
                message="Pipeline completed successfully",
                data={
                    "total_steps": len(self.steps),
                    "completed_steps": sum(
                        1 for s in self.steps if s.status == StepStatus.COMPLETED
                    ),
                    "skipped_steps": sum(
                        1 for s in self.steps if s.status == StepStatus.SKIPPED
                    ),
                },
            )

        except Exception as e:
            logger.exception("Pipeline execution failed")
            self._report_progress(0, f"Pipeline failed: {str(e)}")
            return StepResult(
                success=False,
                message=f"Pipeline failed: {str(e)}",
            )

    def _restore_from_checkpoint(self, state: Dict[str, Any]):
        """
        从断点恢复

        Args:
            state: 断点状态
        """
        logger.info("Restoring from checkpoint")

        steps_state = state.get("steps", {})
        current_step = state.get("current_step")

        for i, step in enumerate(self.steps):
            if step.name in steps_state:
                step_state = steps_state[step.name]
                step.status = StepStatus(step_state["status"])
                step.start_time = None
                step.end_time = None
                step.error_message = step_state.get("error_message", "")

            # 找到当前步骤
            if current_step and step.name == current_step:
                self.current_step_index = i
                logger.info(f"Resuming from step: {step.name}")
                break
            # 如果当前步骤已完成，跳到下一步
            elif step.status == StepStatus.COMPLETED:
                self.current_step_index = i + 1

    def get_status(self) -> Dict[str, Any]:
        """
        获取流水线状态

        Returns:
            状态字典
        """
        return {
            "name": self.name,
            "total_steps": len(self.steps),
            "current_step_index": self.current_step_index,
            "current_step": (
                self.steps[self.current_step_index].name
                if self.current_step_index < len(self.steps)
                else None
            ),
            "steps": [step.get_status_dict() for step in self.steps],
            "has_checkpoint": self.checkpoint_manager.has_checkpoint(),
        }

    def run_step(self, step_name: str, force: bool = False) -> StepResult:
        """
        运行指定步骤

        Args:
            step_name: 步骤名称
            force: 是否强制执行

        Returns:
            执行结果
        """
        # 找到步骤
        step = next((s for s in self.steps if s.name == step_name), None)
        if not step:
            return StepResult(
                success=False,
                message=f"Step not found: {step_name}",
            )

        # 设置强制执行
        if force:
            step.force_run = True
            step.status = StepStatus.PENDING

        return step.execute()

    def skip_step(self, step_name: str):
        """
        跳过指定步骤

        Args:
            step_name: 步骤名称
        """
        step = next((s for s in self.steps if s.name == step_name), None)
        if step:
            step.status = StepStatus.SKIPPED
            logger.info(f"Step skipped: {step_name}")

    def reset(self):
        """重置流水线"""
        for step in self.steps:
            step.status = StepStatus.PENDING
            step.start_time = None
            step.end_time = None
            step.error_message = ""

        self.current_step_index = 0
        logger.info("Pipeline reset")

    def clear_checkpoint(self):
        """清理断点"""
        self.checkpoint_manager.clear_checkpoint(self.name)
        logger.info("Checkpoint cleared")
